import os
import time
import copy
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
from PIL import Image
import argparse

# 设置中文字体支持
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]


# 数据转换
def get_transforms():
    """定义训练和验证的数据转换"""
    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
        'val': transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
        'test': transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
    }
    return data_transforms


# 加载数据集
def load_datasets(data_dir):
    """加载训练集、验证集和测试集"""
    data_transforms = get_transforms()

    image_datasets = {
        'train': datasets.ImageFolder(os.path.join(data_dir, 'train'), data_transforms['train']),
        'val': datasets.ImageFolder(os.path.join(data_dir, 'val'), data_transforms['val']),
        'test': datasets.ImageFolder(os.path.join(data_dir, 'test'), data_transforms['test'])
    }

    dataloaders = {
        x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4, shuffle=True, num_workers=4)
        for x in ['train', 'val', 'test']
    }

    dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val', 'test']}
    class_names = image_datasets['train'].classes

    return dataloaders, dataset_sizes, class_names


# 初始化模型
def initialize_model(num_classes, feature_extract=True):
    """初始化ResNet18模型"""
    model_ft = models.resnet18(pretrained=True)

    # 设置是否只训练最后一层
    if feature_extract:
        for param in model_ft.parameters():
            param.requires_grad = False

    # 修改最后一层以适应分类数量
    num_ftrs = model_ft.fc.in_features
    model_ft.fc = nn.Linear(num_ftrs, num_classes)

    return model_ft


# 训练模型
def train_model(model, dataloaders, criterion, optimizer, scheduler, dataset_sizes, num_epochs=25):
    """训练模型并返回最佳模型"""
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print(f'Epoch {epoch}/{num_epochs - 1}')
        print('-' * 10)

        # 每个epoch都有一个训练和验证阶段
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # 训练模式
            else:
                model.eval()  # 评估模式

            running_loss = 0.0
            running_corrects = 0

            # 迭代数据
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # 零梯度
                optimizer.zero_grad()

                # 前向传播
                # 只有在训练时才跟踪历史
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # 只有在训练阶段才进行反向传播和优化
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # 统计
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            if phase == 'train':
                scheduler.step()

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')

            # 深拷贝模型
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())

        print()

    time_elapsed = time.time() - since
    print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
    print(f'Best val Acc: {best_acc:4f}')

    # 加载最佳模型权重
    model.load_state_dict(best_model_wts)
    return model


# 评估模型
def evaluate_model(model, dataloaders, dataset_sizes, class_names):
    """评估模型并显示混淆矩阵"""
    model.eval()
    class_correct = list(0. for i in range(len(class_names)))
    class_total = list(0. for i in range(len(class_names)))
    confusion_matrix = np.zeros((len(class_names), len(class_names)))

    with torch.no_grad():
        for inputs, labels in dataloaders['test']:
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            c = (predicted == labels).squeeze()

            for i in range(len(labels)):
                label = labels[i]
                class_correct[label] += c[i].item()
                class_total[label] += 1
                confusion_matrix[label][predicted[i]] += 1

    # 显示每个类别的准确率
    for i in range(len(class_names)):
        print(f'Accuracy of {class_names[i]}: {100 * class_correct[i] / class_total[i]:2.2f}%')

    # 计算总体准确率
    overall_accuracy = 100 * sum(class_correct) / sum(class_total)
    print(f'Overall accuracy: {overall_accuracy:2.2f}%')

    # 绘制混淆矩阵
    plt.figure(figsize=(10, 8))
    plt.imshow(confusion_matrix, interpolation='nearest', cmap=plt.cm.Blues)
    plt.title('Confusion matrix')
    plt.colorbar()
    tick_marks = np.arange(len(class_names))
    plt.xticks(tick_marks, class_names, rotation=45)
    plt.yticks(tick_marks, class_names)

    # 在混淆矩阵上标注数值
    thresh = confusion_matrix.max() / 2.
    for i in range(confusion_matrix.shape[0]):
        for j in range(confusion_matrix.shape[1]):
            plt.text(j, i, format(int(confusion_matrix[i, j]), 'd'),
                     horizontalalignment="center",
                     color="white" if confusion_matrix[i, j] > thresh else "black")

    plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label')
    plt.savefig('confusion_matrix.png')
    plt.close()

    return overall_accuracy


# 预测单张图像
def predict_image(model, image_path, class_names):
    """预测单张图像的类别"""
    # 加载图像并进行预处理
    image = Image.open(image_path).convert('RGB')
    transform = get_transforms()['test']
    image_tensor = transform(image).unsqueeze(0)
    image_tensor = image_tensor.to(device)

    # 预测
    model.eval()
    with torch.no_grad():
        outputs = model(image_tensor)
        probs = torch.nn.functional.softmax(outputs, dim=1)
        conf, preds = torch.max(probs, 1)

    # 显示结果
    plt.figure(figsize=(6, 6))
    plt.imshow(image)
    plt.title(f'预测结果: {class_names[preds.item()]} ({conf.item():.2%})')
    plt.axis('off')
    plt.tight_layout()
    plt.savefig('prediction_result.png')
    plt.close()

    print(f'预测类别: {class_names[preds.item()]}')
    print(f'置信度: {conf.item():.2%}')

    # 返回前3个预测结果
    top3_probs, top3_preds = torch.topk(probs, 3)
    print("\nTop 3 预测结果:")
    for i in range(3):
        print(f"{i + 1}. {class_names[top3_preds[0][i]]}: {top3_probs[0][i]:.2%}")

    return class_names[preds.item()], conf.item()


# 主函数
def main():
    parser = argparse.ArgumentParser(description='ResNet18图像分类器')
    parser.add_argument('--data_dir', type=str, default='data', help='数据集目录')
    parser.add_argument('--num_epochs', type=int, default=25, help='训练轮数')
    parser.add_argument('--feature_extract', action='store_true', help='是否只训练最后一层')
    parser.add_argument('--model_path', type=str, default='resnet18_model.pth', help='模型保存路径')
    parser.add_argument('--predict', type=str, help='预测单张图像')

    args = parser.parse_args()

    global device
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    if args.predict:
        # 预测模式
        print("加载模型...")
        dataloaders, dataset_sizes, class_names = load_datasets(args.data_dir)
        num_classes = len(class_names)
        model = initialize_model(num_classes, feature_extract=False)
        model = model.to(device)
        model.load_state_dict(torch.load(args.model_path))

        print(f"对图像 {args.predict} 进行预测")
        predict_image(model, args.predict, class_names)

    else:
        # 训练模式
        print("加载数据...")
        dataloaders, dataset_sizes, class_names = load_datasets(args.data_dir)
        num_classes = len(class_names)
        print(f"发现 {num_classes} 个类别: {class_names}")

        # 初始化模型
        print("初始化模型...")
        model_ft = initialize_model(num_classes, args.feature_extract)
        model_ft = model_ft.to(device)

        # 定义损失函数和优化器
        criterion = nn.CrossEntropyLoss()

        # 只更新需要训练的参数
        params_to_update = model_ft.parameters()
        print("需要训练的参数:")
        if args.feature_extract:
            params_to_update = []
            for name, param in model_ft.named_parameters():
                if param.requires_grad:
                    params_to_update.append(param)
                    print("\t", name)
        else:
            for name, param in model_ft.named_parameters():
                if param.requires_grad:
                    print("\t", name)

        optimizer_ft = optim.SGD(params_to_update, lr=0.001, momentum=0.9)

        # 学习率调度器
        exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

        # 训练模型
        print("开始训练模型...")
        model_ft = train_model(model_ft, dataloaders, criterion, optimizer_ft,
                               exp_lr_scheduler, dataset_sizes, num_epochs=args.num_epochs)

        # 保存模型
        torch.save(model_ft.state_dict(), args.model_path)
        print(f"模型已保存至 {args.model_path}")

        # 评估模型
        print("评估模型...")
        evaluate_model(model_ft, dataloaders, dataset_sizes, class_names)


if __name__ == "__main__":
    main()